<?php

namespace Mcp\Server\Transport;

use Mcp\Server\TransportInterface;
use Psr\Log\LoggerInterface;
use Psr\Log\NullLogger;
use Swoole\Coroutine\WebSocket\Frame;
use Swoole\WebSocket\Server;
use Symfony\Component\Uid\Uuid;

class SwooleWebSocketTransport implements TransportInterface
{
    private ?Server $server = null;
    private $messageListener;
    private $sessionEndListener;
    private array $connections = []; // 存储连接与会话ID的映射

    public function __construct(
        private string $host = '0.0.0.0',
        private int $port = 9501,
        private LoggerInterface $logger = new NullLogger()
    ) {}

    public function initialize(): void
    {
        // 初始化 Swoole WebSocket 服务器（协程模式）
        $this->server = new Server($this->host, $this->port);
        $this->server->set([
            'enable_coroutine' => true, // 启用协程
            'worker_num' => 4,
        ]);

        // 注册连接事件
        $this->server->on('open', function (Server $server, $request) {
            $sessionId = Uuid::v4(); // 为新连接创建会话ID
            $this->connections[$request->fd] = $sessionId;
            $this->logger->info('WebSocket connection opened', [
                'fd' => $request->fd,
                'session_id' => $sessionId->toRfc4122()
            ]);
        });

        // 注册消息事件（触发 SDK 的 message 回调）
        $this->server->on('message', function (Server $server, Frame $frame) {
            $fd = $frame->fd;
            $sessionId = $this->connections[$fd] ?? null;
            if ($this->messageListener && $sessionId) {
                // 调用 SDK 注册的消息处理逻辑
                call_user_func($this->messageListener, $frame->data, $sessionId);
            }
        });

        // 注册关闭事件（触发会话销毁）
        $this->server->on('close', function (Server $server, int $fd) {
            $sessionId = $this->connections[$fd] ?? null;
            if ($this->sessionEndListener && $sessionId) {
                call_user_func($this->sessionEndListener, $sessionId);
            }
            unset($this->connections[$fd]);
            $this->logger->info('WebSocket connection closed', ['fd' => $fd]);
        });
    }

    public function onMessage(callable $listener): void
    {
        $this->messageListener = $listener; // 接收 SDK 注册的消息处理器
    }

    public function listen(): mixed
    {
        $this->logger->info('Starting Swoole WebSocket server', [
            'host' => $this->host,
            'port' => $this->port
        ]);
        $this->server->start(); // 启动服务器（阻塞当前进程）
        return null;
    }

    public function send(string $data, array $context): void
    {
        // 从上下文获取目标连接 FD（需在消息处理时传递）
        if (isset($context['fd']) && isset($this->server)) {
            $this->server->push($context['fd'], $data);
            $this->logger->debug('Sent data via WebSocket', [
                'fd' => $context['fd'],
                'data' => $data
            ]);
        }
    }

    public function onSessionEnd(callable $listener): void
    {
        $this->sessionEndListener = $listener; // 接收 SDK 注册的会话销毁处理器
    }

    public function close(): void
    {
        if ($this->server) {
            $this->server->shutdown();
        }
    }
}